import numpy as np
#扩散算符
def diff_opr1(myfun,nu,dx,dy):
    diff_y = nu * (myfun[2:, 1:-1] + myfun[:-2, 1:-1] - 2 * myfun[1:-1, 1:-1]) / dy / dy
    diff_x = nu * (myfun[1:-1, 2:] + myfun[1:-1, :-2] - 2 * myfun[1:-1, 1:-1]) / dx / dx
    return diff_x+diff_y

def diff_opr2(myfun,nu,dx,dy):
    diff_y = nu * (myfun[2:, 1:-1] + myfun[:-2, 1:-1] - 2 * myfun[1:-1, 1:-1]) / dy / dy
    diff_x = nu * (myfun[1:-1, 2:] + myfun[1:-1, :-2] - 2 * myfun[1:-1, 1:-1]) / dx / dx
    return diff_x,diff_y

'''计算流函数，本质上是解泊松方程。要注意边界条件的处理'''
def compute_stream_xperiod(omega_int,Nx,Ny, dx, dy, tol=1e-5, psi_init=None,max_iters=100000,min_iters=500):
    """
    Jacobi iteration,x period,y diriclet
    主输入:涡度场值omega_int
    辅输入:dx,dy,网格长,用于算差分矩阵
    主返回:计算得到的psi_init
    """
    if psi_init is None:
        psi_init = np.zeros((Ny, Nx))
    iter = 0
    flag_need_more=True
    """进入迭代
    判断条件：
    1.迭代次数小于min_iters,一定要继续迭代
    2.迭代次数大于min_iters,小于max_iters时,通过flag_need_more 决定是否要继续迭代,如果为否,停止迭代
    3.迭代次数大于max_iters,一定要停止迭代
    flag_need_more判断逻辑:

    """
    while (flag_need_more and iter<max_iters) or iter<min_iters:
        psi_new = psi_init.copy()
        #对中间区域格点的迭代过程，这里把差分矩阵在雅克比迭代情况的表达式写进来了。计算采用矩阵加法）
        psi_new[1:-1, 1:-1] = (omega_int \
                               + (psi_init[2:, 1:-1] + psi_init[:-2, 1:-1]) / dy / dy \
                               + (psi_init[1:-1, 2:] + psi_init[1:-1, :-2]) / dx / dx) / (2 / dx / dx + 2 / dy / dy)
        psi_new[-1, :]=0
        psi_new[0, :]=0  #底部
        psi_new[1:-1,0]=psi_new[1:-1,-2]    #周期条件：左侧边界点用右侧内点赋值
        psi_new[1:-1,-1]=psi_new[1:-1,1]    #周期条件：右侧边界点用左侧内点赋值
        
        
        psi_init = psi_new
        """计算迭代后结果是否对"""
        if iter%500==0:
            diff_x,diff_y=diff_opr2(psi_init,1,dx,dy)
            omega_numeric_tmp=-diff_x-diff_y
            error_tmp=omega_int-omega_numeric_tmp
            errm=np.max(np.abs(error_tmp))
            omega_ave=np.average(omega_int)

            if errm<tol:
                flag_need_more=False
            if omega_ave>1:
                if errm/omega_ave<tol:
                    flag_need_more=False

        iter += 1   
        
    end_signal = True if iter == 1 else False  #只用迭代1次时，end_signal返回真，这将用于停止迭代。因此此程序似乎在求从零直到稳态的过程。
    if iter>min_iters:
        print(f'jaccobi iteration:{iter},error omega max:{errm}')

    return psi_init, end_signal

'''计算流函数，本质上是解泊松方程。要注意边界条件的处理'''
def compute_stream_diriclet(omega_int,Nx,Ny, dx, dy, tol=1e-7, psi_init=None,max_iters=100000,min_iters=1000):
    """
    Jacobi iteration
    主输入:涡度场值omega_int
    辅输入:dx,dy,网格长,用于算差分矩阵
    主返回:计算得到的psi_init
    """
    if psi_init is None:
        psi_init = np.zeros((Ny, Nx))
    iter = 0
    flag_need_more=True
    """进入迭代
    判断条件：
    1.迭代次数小于min_iters,一定要继续迭代
    2.迭代次数大于min_iters,小于max_iters时,通过flag_need_more 决定是否要继续迭代,如果为否,停止迭代
    3.迭代次数大于max_iters,一定要停止迭代
    flag_need_more判断逻辑:

    """
    while (flag_need_more and iter<max_iters) or iter<min_iters:
        psi_new = psi_init.copy()
        #对中间区域格点的迭代过程，这里把差分矩阵在雅克比迭代情况的表达式写进来了。计算采用矩阵加法）
        psi_new[1:-1, 1:-1] = (omega_int \
                               + (psi_init[2:, 1:-1] + psi_init[:-2, 1:-1]) / dy / dy \
                               + (psi_init[1:-1, 2:] + psi_init[1:-1, :-2]) / dx / dx) / (2 / dx / dx + 2 / dy / dy)
        '''
        如果边界上流函数一直为零,初始时设成0就好,后面不动它
        顶盖第二类边条件(收敛性存在问题):psi_new[-1, 1:-1]=psi_new[-2, 1:-1]+dy
        '''
        #psi_new[-1, :]=psi_new[-2, :]   #顶部
        psi_new[-1, :]=0
        psi_new[0, :]=0  #底部
        psi_new[1:-1,0]=0
        psi_new[1:-1,-1]=0
        
        
        psi_init = psi_new
        """计算迭代后结果是否对"""
        if iter%500==0:
            diff_x,diff_y=diff_opr2(psi_init,1,dx,dy)
            omega_numeric_tmp=-diff_x-diff_y
            error_tmp=omega_int-omega_numeric_tmp
            errm=np.max(np.abs(error_tmp))
            omega_ave=np.average(omega_int)
            #change = np.max(np.abs(psi_new - psi_init))  #change为迭代后矩阵元素变化最大值
            #print(errm)
            if errm<tol:
                flag_need_more=False
            if omega_ave>1:
                if errm/omega_ave<tol:
                    flag_need_more=False

        iter += 1   
        
    end_signal = True if iter == 1 else False  #只用迭代1次时，end_signal返回真，这将用于停止迭代。因此此程序似乎在求从零直到稳态的过程。
    if iter>min_iters:
        print(f'jaccobi iteration:{iter},error omega max:{errm}')

    return psi_init, end_signal



def apply_bnd_omega_yfreeslip_xperiod(psi, u, v, omega):
    #左右边界周期条件
    omega[:,0]=omega[:,-2]
    omega[:,-1]=omega[:,1]

    #顶部底部涡量直接设为零
    omega[0, 1:-1] =0
    omega[-1, 1:-1] =0

def apply_bnd_omega_solidwall(psi, u, v, omega,dx,dy):
    #原来的左右壁面涡量条件
    omega[:, 0] = 2 / dx / dx * (psi[:, 0] - psi[:, 1]) - 2 / dx * v[:, 0]
    omega[:, -1] = 2 / dx / dx * (psi[:, -1] - psi[:, -2]) + 2 / dx * v[:, -1]
    
    #如果底部仍用壁面边条件
    omega[0, 1:-1] = 2 / dy / dy * (psi[0, 1:-1] - psi[1, 1:-1]) + 2 / dy * u[0, 1:-1]
    
    #顶部采用的壁面涡量条件，这是顶部切向速度要求
    omega[-1, 1:-1] = 2 / dy / dy * (psi[-1, 1:-1] - psi[-2, 1:-1]) - 2 / dy * u[-1,1 :-1]